import base64
import json
import copy

from flask import Flask, request, jsonify, send_from_directory, render_template, abort
import requests
import threading
import time

from flask_cors import CORS
from flask_socketio import SocketIO, emit
from flask_talisman import Talisman

import agg_config
import urllib3
from sql_options import *
from agg import ServerAgg
from utils.options import args_parser, get_config
import torch
from enum import Enum

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

app = Flask(__name__, static_folder="./dist/static/",
            template_folder="./dist")
cors_origins = os.getenv('CORS_ORIGIN', '10.3.28.134')
# CORS(app, resources={r"/*": {"origins": cors_origins}})  # 允许所有域访问所有路由
CORS(app, resources={r"/*": {"origins": "*"}})  # 允许所有域访问所有路由
csp = {
    'default-src': ["'self'"],
    'script-src': ["'self'", "https://buttons.github.io/", "https://cdnjs.cloudflare.com",
                   "https://kit.fontawesome.com", "'nonce-randomNonceValue'"],
    'style-src': ["'self'", "https://fonts.googleapis.com", "https://cdnjs.cloudflare.com", "'nonce-randomNonceValue'"],
    'img-src': ["'self'", "data:"],
    'font-src': ["'self'", "https://fonts.gstatic.com", "https://cdnjs.cloudflare.com", "https://ka-f.fontawesome.com",
                 "data:"],
    # 'frame-ancestors': ["'none'"],
    'connect-src': ["'self'", "https://ka-f.fontawesome.com"],
    'form-action': ["'self'"],
    'object-src': ["'none'"],
    'base-uri': ["'self'"],
}

# talisman = Talisman(app, content_security_policy=csp, force_https=False)
# socketio = SocketIO(app, cors_allowed_origins="*")

args = get_config()
# conn = create_connection()
# if conn:
#     create_agg_data_table(conn)

server = None

class ServerStep(Enum):
    NONE = "step_none"  # 未开始
    INIT_DONE = "step_init_done"  # 初始化完成
    AGG_DONE = "step_agg_done"  # 聚合完成
    POST_DONE = "step_post_done"  # 发布完成

    # 添加一个方法，返回对应的中文描述
    def get_description(self):
        descriptions = {
            ServerStep.NONE: "未开始",
            ServerStep.INIT_DONE: "初始化完成",
            ServerStep.AGG_DONE: "聚合完成",
            ServerStep.POST_DONE: "发布完成",
        }
        return descriptions.get(self, "未知步骤")

def get_default_data():
    return {
        "states": [False, False, "server information"],
        "clients": {},
        "step_status": ServerStep.NONE.value
    }


cache_data = get_default_data()


@app.route("/init", methods=['GET', 'POST'])
def init():
    # conn = create_connection()
    # if conn:
    #     create_agg_data_table(conn)
    # parse args
    # 连接数据库
    conn = create_connection()
    global server
    # 更新多个值
    args.update({
        'device': torch.device('cuda:0' if torch.cuda.is_available() and 0 != -1 else 'cpu'),
        'method': 'VeriTrac',
        'global_ep': 1
    })
    # args = args_parser()
    # args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
    # args.method = 'VeriTrac'
    # args.global_ep = 1

    response_trace = requests.post(
        f'{agg_config.SERVER_TRACE_HOST}/get_server_params', verify=False
    )
    data = response_trace.json()
    print("agg init start")
    params = data.get('params')
    server = ServerAgg(args, params)
    print("agg init end")

    # 前端初始化
    clients_simpledata, cache = init_data(conn)
    cache_data["step_status"] = ServerStep.INIT_DONE.value
    return jsonify({'clients_simpledata': clients_simpledata})

@app.route("/reset", methods=['POST'])
def reset():
    global server, cache_data
    server = None
    cache_data = get_default_data()
    return jsonify({'status': 'success'}), 200

def init_data(conn=None):
    """
    初始化数据,并返回缓存和客户端列表
    return: 客户端数据, 缓存数据
    """
    global cache_data
    if not conn:
        # 连接数据库
        conn = create_connection()
    # 前端初始化
    set_id(conn)
    # id, client_id,userid,username
    clients_simpledata = get_simpledata_clients_data(conn)
    conn.close()
    client_map = {}
    cache_data = get_default_data()
    for item in clients_simpledata:
        client = {
            "name": item["username"],
            "client_id": item["client_id"],
            "status": False,
            "userid": item["userid"],
            "info": ""
        }
        client_map[item["client_id"]] = client
    cache_data["clients"] = client_map
    return clients_simpledata, cache_data


@app.route('/upload', methods=['POST'])
def upload():  # 连接数据库
    conn = create_connection()
    data = request.json
    upload_data = data.get('upload_data', {})
    info = upload_data.get('local_BFV')
    if 'local_BFV' in upload_data and upload_data['local_BFV'] is not None:
        upload_data['local_BFV'] = base64.b64decode(upload_data['local_BFV'])
    local_BFV = upload_data.get('local_BFV')
    local_sign1 = upload_data.get('local_sign1')
    local_sign2 = upload_data.get('local_sign2')
    pk_sign_1 = upload_data.get('pk_sign_1')
    pk_sign_2 = upload_data.get('pk_sign_2')
    client_id = upload_data.get('client_id')
    client_num = upload_data.get('client_num')
    insert_or_update_agg_data(conn, client_id, client_sign1=local_sign1, client_sign2=local_sign2,
                              client_BFV=local_BFV, pk_sign_1=pk_sign_1, pk_sign_2=pk_sign_2, client_num=client_num)
    conn.close()
    client_map = cache_data.get("clients")
    if client_map is not None and client_map.get(client_id):
        client = client_map.get(client_id)
        client["info"] = info
        client["status"] = True
    # socketio.emit('update_client', {'client_id': client_id, 'status': True, "info": info})
    return jsonify({'status': 'success'}), 200


@app.route('/aggregate', methods=['POST'])
def aggregate():
    # 连接数据库
    conn = create_connection()
    data = request.json
    global verify_aggregation, global_BFV_bytes, pk_sign_1_s, sig1_s_recover
    verify_aggregation, global_BFV_bytes, pk_sign_1_s, sig1_s_recover = server.aggregation(
        server_state=data.get('server_state'), connection=conn)
    print('data', verify_aggregation)
    if verify_aggregation == 0:
        all_clients = get_basedata_clients_data(conn)
        for client in all_clients:
            client_id = client[0]
            client_ip = client[1]
            # client_port = client[2]
            headers = {'Content-Type': 'application/json', 'Client-ID': client_id}
            response = requests.post(
                f'{client_ip}/re_train', headers=headers, verify=False)
            print(f"Sent data to client: {response.json()}")
    conn.close()
    cache_data["states"][0] = verify_aggregation != 0
    cache_data["states"][1] = data.get('server_state')
    cache_data["states"][2] = global_BFV_bytes
    if verify_aggregation == 1:
        cache_data["step_status"] = ServerStep.AGG_DONE.value
    else:
        init_data()
        cache_data["step_status"] = ServerStep.INIT_DONE.value
    return jsonify({'verify_aggregation': verify_aggregation, 'aggregation_model': global_BFV_bytes}), 200
    # return jsonify({'status': 'success'}), 200


@app.route('/post_agg_result', methods=['POST'])
def post_agg_result():
    # 连接数据库
    conn = create_connection()
    # 获取所有客户端基本信息
    all_clients = get_basedata_clients_data(conn)
    conn.close()
    print("All Clients Data:", all_clients)
    resp_state = []
    for client in all_clients:
        client_id = client[0]
        client_ip = client[1]
        # client_port = client[2]
        print(f"Client ID: {client_id}, IP: {client_ip}")
        headers = {'Content-Type': 'application/json', 'Client-ID': client_id}
        data = {
            'global_BFV_bytes': global_BFV_bytes,
            'pk_sign_1_s': pk_sign_1_s,
            'sig1_s': sig1_s_recover,
        }
        response = requests.post(
            f'{client_ip}/update',
            json=data, headers=headers, verify=False)
        if response.status_code == 200:
            print(f"Sent data to client: {response.json()}")
            resp_state.append(True)
        else:
            print("Sent data to client error! {}".format(response.text))
            resp_state.append(False)
    init_data()
    cache_data["step_status"] = ServerStep.POST_DONE.value
    return jsonify({'status': 'success'}), 200


@app.route('/post_clients_data', methods=['GET'])
def post_clients_data():
    # 连接数据库
    conn = create_connection()
    all_clients = get_all_clients_data(conn)
    conn.close()
    print("All Clients Data:", type(all_clients), len(all_clients))
    for client in all_clients:
        print(type(client.get('client_BFV')))
        if 'client_BFV' in client and client['client_BFV'] is not None:
            client['client_BFV'] = base64.b64encode(client['client_BFV']).decode('utf-8')

        # 计算数据包大小
    response_size = len(json.dumps({'all_clients': all_clients}).encode('utf-8'))

    print(f"Response size: {response_size} bytes")  # 输出数据包大小
    return jsonify({'all_clients': all_clients}), 200


@app.route('/delete_client', methods=['DELETE'])
def delete_client():
    client_id = request.json.get('client_id')
    print(f"Client ID: {client_id}")
    if not client_id:
        return jsonify({'error': 'client_id is required'}), 400
    deleted_rows = delete_client_data(client_id)
    print(deleted_rows)
    if deleted_rows is None:
        return jsonify({'error': 'An error occurred while deleting the data'}), 500
    elif deleted_rows > 0:
        return jsonify({'message': 'Client data deleted successfully'}), 200
    else:
        return jsonify({'error': 'Client ID not found'}), 404


@app.route('/data', methods=['POST'])
def receive_data():
    # 连接数据库
    conn = create_connection()
    data = request.json
    print(f"Received data: {data}")
    client_id = request.headers.get('Client-ID')
    client_ip = data.get('client_ip')
    # client_port = data.get('client_port')
    userid = data.get('userid')
    username = data.get('username')
    print(f"Received data from client {client_id}: {data}")
    insert_or_update_agg_data(conn, client_id=client_id, client_ip=client_ip, userid=userid,
                              username=username)

    conn.close()
    return jsonify({'status': 'success'}), 200


@app.route('/ping', methods=['GET'])
def ping():
    return jsonify({'status': 'alive'}), 200


def run_flask():
    # app.run(host='0.0.0.0', port=5002, ssl_context=(
    #     f'{Path(__file__).parent}/certs/server.crt',
    #     f'{Path(__file__).parent}/certs/server.key'
    # ))
    app.config['DEBUG'] = False
    app.run(host='0.0.0.0', port=5002)
    # socketio.run(app, host='0.0.0.0', port=5002, allow_unsafe_werkzeug=True)


# WebSocket 连接事件
# @socketio.on('connect')
# def handle_connect():
#     print('Client connected')
#
#
# @socketio.on('disconnect')
# def handle_disconnect():
#     print('Client disconnected')


def connect_to_trace():
    while True:
        try:
            response = requests.get(f'{agg_config.SERVER_TRACE_HOST}/ping', verify=False)
            print(f"Connected to trace server: {response.text}")
            # init()

            break
        except requests.exceptions.RequestException:
            print("Connection failed, retrying...")
            time.sleep(5)


# 路由处理 favicon.ico 请求
@app.route('/favicon.ico')
def favicon():
    return send_from_directory(os.path.join(app.root_path, 'static'), 'favicon.png', mimetype='image/png')


@app.route('/')
@app.route('/<path:path>')
def serve(path='index.html'):
    if path != '' and os.path.exists(os.path.join(app.template_folder, path)):
        # 检查路径是否指向模板文件（如 .html 文件）
        if path.endswith('.html'):
            return render_template(path)
        else:
            # 如果路径指向静态文件，则返回静态文件
            return send_from_directory(app.template_folder, path)
    else:
        return render_template('index.html')

def get_trace_step_status():
    response = requests.get(f'{agg_config.SERVER_TRACE_HOST}/trace_status?step_status=true', verify=False)
    if response.status_code == 200:
        step_status = response.json()["step_status"]
        print("Get trace step status successfully: {}".format(step_status))
        return step_status
    else:
        print("Get trace step status error: {}".format(response.text))
        return None

@app.route('/agg_status', methods=['GET'])
def agg_status():
    """
    返回 server agg 状态及会话缓存
    """
    step_status = request.args.get('step_status')
    if step_status:
        return {"step_status": cache_data["step_status"]}
    trace_step = get_trace_step_status()
    result = cache_data.copy()
    result["clients"] = list(result["clients"].values())
    result["trace_step_status"] = trace_step
    return result


# 定义禁止访问的主机列表
BLOCKED_HOSTS = ['169.254.169.254', 'aws.zaproxy.org', '100.100.100.200', 'alibaba.zaproxy.org']


@app.before_request
def block_metadata_access():
    if request.host in BLOCKED_HOSTS:
        abort(403)  # 禁止访问


if __name__ == "__main__":
    # 检查和初始化数据库
    conn = create_connection()
    if conn:
        create_agg_data_table(conn)
    conn.close()
    # flask_thread = threading.Thread(target=run_flask)
    # flask_thread.start()

    connect_thread = threading.Thread(target=connect_to_trace)
    connect_thread.start()

    # Wait for threads to complete
    connect_thread.join()
    run_flask()
    # flask_thread.join()

    # conn = create_connection()
    # if conn:
    #     create_agg_data_table(conn)
    #
    # app.config['DEBUG'] = False
    # # app.run(host='0.0.0.0', port=5003)
    # socketio.run(app, host='0.0.0.0', port=5002)
    # conn.close()
